import stix2
from pycti import StixCoreRelationship

from . import utils


def create_stix_malware(connector, malware):
    return stix2.Malware(
        id=malware["id"],
        is_family=True,
        name=utils.sanitizer("name", malware),
        description=utils.sanitizer("description", malware),
        modified=utils.sanitizer("last_updated", malware),
        aliases=utils.clean_intrusionset_aliases(malware),
        confidence=connector.helper.connect_confidence_level,
        created_by_ref=connector.identity["standard_id"],
        object_marking_refs=[stix2.TLP_AMBER.get("id")],
    )


def create_stix_relationship(
    connector, rel_type, source, target, start_time, stop_time
):
    return stix2.Relationship(
        id=StixCoreRelationship.generate_id(
            rel_type, source, target, start_time, stop_time
        ),
        relationship_type=rel_type,
        source_ref=source,
        target_ref=target,
        start_time=start_time,
        stop_time=stop_time,
        allow_custom=True,
        created_by_ref=connector.identity["standard_id"],
    )


def create_stix_industry(connector, base_object, industry):
    stix_identity = stix2.Identity(
        id=industry["id"],
        created_by_ref=connector.identity["standard_id"],
        name=industry["name"],
        identity_class="class",
        allow_custom=True,
    )
    stix_relationship = create_stix_relationship(
        connector,
        "targets",
        base_object.get("id"),
        stix_identity.get("id"),
        industry.get("first_seen"),
        industry.get("last_seen", None),
    )
    return [stix_identity, stix_relationship]


def create_stix_vulnerability(connector, stix_intrusionset, cve):
    stix_vulnerability = stix2.Vulnerability(
        id=cve["id"],
        created_by_ref=connector.identity["standard_id"],
        name=cve["cve_id"],
        allow_custom=True,
    )
    stix_relationship = create_stix_relationship(
        connector,
        "targets",
        stix_intrusionset.get("id"),
        stix_vulnerability.get("id"),
        cve.get("first_seen"),
        cve.get("last_seen", None),
    )
    return [stix_vulnerability, stix_relationship]


def process(connector, malware):
    malware_id = malware.get("id")

    connector.helper.log_debug(f"Processing malware {malware_id} ...")

    stix_malware = create_stix_malware(connector, malware)
    malware_details = connector.api.malware(malware_id)

    items = [stix_malware]

    for industry in malware_details.get("industries", []):
        items += create_stix_industry(connector, stix_malware, industry)

    for cve in malware_details.get("cve", []):
        items += create_stix_vulnerability(connector, stix_malware, cve)

    bundle = stix2.Bundle(objects=items, allow_custom=True)

    if bundle is None:
        connector.helper.log_error(
            f"Could not process malware {malware_id}. Skipping ..."
        )

    return bundle
